#!/usr/bin/env python
'''
    Kill the process who eats memory fastest & whose memory usage is the most
    while the total memory usage is `gt' a certain value.
    If anyone ask me why does it not work well? I don't know either...
'''
__author__ = 'huoxy'
__email__ = 'sykp241095@gmail.com'

import os
import sys
import time
import logging
from optparse import OptionParser, OptParseError

# options, from command
options = {}
args = ()

# vars, from system or os
memtotal = 0

# global vars
count = 0
percent = 0
logfile = ''
interval = 0
rate = 0.0

# ps infos
old_psinfos = {}

def daemonize(stdin = '/dev/null', stdout = '/dev/null', stderr = '/dev/null'):
    try:
        pid = os.fork()
        if pid > 0:
            sys.exit(0) # exit first parent
    except OSError, e:
        errmsg = 'fork #1 failed: (%d) %s' % (e.errno, e.strerror)
        sys.stderr.write(errmsg + '\n')
        logging.error(errmsg)
        sys.exit(1)
        
#    os.chdir('/')
    os.umask(0)
    os.setsid()
        
    # perform second fork.
    try:
        pid = os.fork()
        if pid > 0: sys.exit(0) # exit second parent
    except OSError, e:
        errmsg = 'fork #2 failed: (%d) %s' % (e.errno, e.strerror)
        sys.stderr.write(errmsg + '\n')
        logging.error(errmsg)
        sys.exit(1)

    for f in sys.stdout, sys.stderr: f.flush()

    si = file(stdin, 'r')
    so = file(stdout, 'a+')
    se = file(stderr, 'a+', 0)

    os.dup2(si.fileno(), sys.stdin.fileno())
    os.dup2(so.fileno(), sys.stdout.fileno())
    os.dup2(se.fileno(), sys.stderr.fileno())



def init_options():
    '''initial `options' and `args'
    '''
    usage = '%prog [options] args'
    parser = OptionParser(usage = usage, description = 'kill the memory eater')
    parser.add_option('-c', '--count', help = 'process count to kill',  default = 1)
    parser.add_option('-p', '--percent', help = 'total usage percent of memory', default = 90)
    parser.add_option('-l', '--logfile', help = 'log file to write  into', \
            default = os.path.basename(sys.argv[0]) + '.log')
    parser.add_option('-i', '--interval', help = 'interval between checks', default = 3)
    parser.add_option('-r', '--rate', help = 'memory usage increasement rate', default = 0.01)
    global options, args
    try:
        options, args = parser.parse_args()
    except OptParseError, e:
        sys.exit(1)

def init_sysargs():
    ''' initial vars from system & os
    '''
    global memgsize
    memtotal = get_mem_total()

def init_globals():
    ''' initial `log' from `options'
    '''
    global count, percent, logfile, interval, rate
    count = options.count
    percent = options.percent
    logfile = options.logfile
    interval = options.interval
    rate = options.rate

def init_logging():
    '''initial `logging' module
    '''
    logging.basicConfig(
            loglevel = logging.ERROR,
            format = '%(asctime)s %(levelname)s: %(message)s',
            filename = logfile,
            filemode = 'a')

def get_mem_total():
    '''use the `free' command, the 2nd line, the 2nd field
    '''
    free = os.popen('free -k', 'r')
    free.readline()
    memgsize = int(free.readline().split()[1])
    return memgsize

def get_mem_usage():
    '''use the `free' command, the 3rd line, the 3rd field
    '''
    free = os.popen('free -k', 'r')
    free.readline() # we do not use the title line
    memusage = int(free.readline().split()[2])
    return memusage

def get_mem_usage_percent():
    return round( float(get_mem_usage()) ) / memtotal * 100

def is_critical():
    '''check if memory usage now is greater than our acceptable `percent'
        which was provided in the command options
    '''
    return get_mem_usage_percent > percent and True or False

def get_psinfos():
    '''use the `ps' command
    '''
    psinfos = {}
    try:
        ps = os.popen(r''' ps -eo pid,rss | sed '1d' | sort -nr -k2 ''')
    except OSError, e:
        logging.error(e)
        sys.exit(1)
    return dict( [ info.split() for info in ps.readlines() ] )

def analyze():
    '''compare the `old_psinfos' and `psinfos' dict, decide whether kill one process or not
    '''
    size = memtotal * rate * interval
    psinfos = get_psinfos()
    # clean the pid in old_`psinfos' but not in `psinfos'
    for pid in old_psinfos.keys():
        if pid not in psinfos:
            del old_psinfos[pid]
    for key in psinfos.keys():
        pid, rss = int(key), int(psinfos[key])
        if pid in old_psinfos:
            old_rss = old_psinfos[pid]
            if old_rss - rss > size:
                # TODO if any, do not blame me...
#                kill(pid, 9)
                print 'kill(pid, 9)'
        # if pid is not in old_psinfos, that means `pid' is a new one.
        # we should store it.
        else:
            old_psinfos[pid] = rss

def kill(pid = None, sig = None):
    '''obviously kill `pid' with signal `sig'
    '''
    psinfos = get_psinfos()
    if not pid:
        return
    elif pid not in psinfos:
        return
    else:
        try:
            os.kill(pid, sig)
        except OSError, e:
            logging.error(e)
            sys.exit(1)

def main():
#    daemonize()
    init_options()
    init_sysargs()
    init_globals()
    init_logging()

    while True:
        if is_critical():
            analyze()
            continue
            time.sleep(interval)

if __name__ == '__main__':
    main()
